"""
手写体识别
"""
import paddle
import numpy as np
import matplotlib.pyplot as plt

import paddle.vision.transforms as T

# 数据的加载和预处理
transform = T.Normalize(mean=[127.5],std=[127.5])
# 训练数据集
train_dataset = paddle.vision.datasets.MNIST(mode='train',transform=transform)
# 评估数据集
eval_dataset = paddle.vision.datasets.MNIST(mode='test',transform=transform)

print('训练集样本量：{}，验证集样本量：{}'.format(len(train_dataset),len(eval_dataset)))


# print('图片：')
# print(type(train_dataset[0][0]))
# print(train_dataset[0][0])
# print('标签：')
# print(type(train_dataset[0][1]))
# print(train_dataset[0][1])

# 可视化展示
# plt.figure()
# plt.imshow(train_dataset[0][0].reshape([28,28]),cmap=plt.cm.binary)
# plt.show()

# 模型网络结构搭建
network = paddle.nn.Sequential(
    paddle.nn.Flatten(), # 拉平，将（28,28） => (784)
    paddle.nn.Linear(784,512), # 隐层：线性变换层
    paddle.nn.ReLU(), # 激活函数
    paddle.nn.Linear(512,10) # 输出层
)

# 模型封装
model = paddle.Model(network)

# 模式可视化
model.summary((1,28,28))

# 配置优化器、损失函数、评估指标
model.prepare(paddle.optimizer.Adam(learning_rate=0.001,parameters=network.parameters()),
              paddle.nn.CrossEntropyLoss(),
              paddle.metric.Accuracy())

# 启动模型全流程训练
model.fit(train_dataset, # 训练数据集
          eval_dataset,  # 评估数据集
          epochs=5,      # 训练的总轮次
          batch_size=64, # 训练使用的批大小
          verbose=1)     # 日志展示形式

# 模型评估，根据prepare接口配置的loss和metric进行返回
result = model.evaluate(eval_dataset,verbose=1)
print(result)

# 保存用于后续继续调优训练的模型
model.save('writeRecognition/mnist')

# 保存用于后续推理部署的模型
# model_2.save('infer/mnist', training=False)

# 进行预测操作
# result = model.predict(eval_dataset)
#
# def show_img(img,predict):
#     plt.figure()
#     plt.title('predict: {}'.format(predict))
#     plt.imshow(img.reshape([28,28]),cmap=plt.cm.binary)
#     plt.show()
#
# # 抽样展示
# indexs = [2,15,38,211]
#
# for idx in indexs:
#     show_img(eval_dataset[idx][0],np.argmax(result[0][idx]))

from paddle.static import InputSpec

# 模型放封装，为了后面保存预测模型，这里传入input参数
model_2 = paddle.Model(network,inputs=[InputSpec(shape=[-1,28,28],dtype='float32',name='image')])

# 加载之前保存的训练阶段的模型
model_2.load('writeRecognition/mnist')

# 模型配置
model_2.prepare(paddle.optimizer.Adam(learning_rate=0.001,parameters=network.parameters()),
                paddle.nn.CrossEntropyLoss(),
                paddle.metric.Accuracy())

model_2.fit(train_dataset,
            eval_dataset,
            epochs=2,
            batch_size=64,
            verbose=1)
model_2.save('infer/mnist',training=False)